package com.hcc.repository.gen.metadata.dialect;

import com.hcc.repository.core.utils.ReflectUtils;
import com.hcc.repository.core.utils.StrUtils;
import com.hcc.repository.gen.config.GlobalConfig;
import com.hcc.repository.gen.metadata.MetadataQuery;
import com.hcc.repository.gen.model.ColumnMetadata;
import com.hcc.repository.gen.model.TableMetadata;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * JdbcMetadataQuery
 *
 * @author hushengjun
 * @date 2023/10/18
 */
public class JdbcMetadataQuery implements MetadataQuery {

    private static final String SQL = "SELECT * FROM %s WHERE 1=2";

    private DataSource dataSource;

    @Override
    public void setDataSource(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    @Override
    public List<TableMetadata> buildTableMetadata(List<String> tableNames, GlobalConfig globalConfig) {
        try (Connection connection = dataSource.getConnection()) {
            DatabaseMetaData databaseMetaData = connection.getMetaData();
            ResultSet tableResultsSet = databaseMetaData.getTables(connection.getCatalog(), null, null, new String[]{});

            List<TableMetadata> tableMetadataList = new ArrayList<>();
            while (tableResultsSet.next()) {
                String tableName = tableResultsSet.getString("TABLE_NAME");
                if (!tableNames.contains(tableName)) {
                    continue;
                }
                TableMetadata tableMetadata = new TableMetadata();
                tableMetadata.setGlobalConfig(globalConfig);
                tableMetadata.setName(tableName);
                String remarks = tableResultsSet.getString("REMARKS");
                tableMetadata.setComment(StrUtils.isEmpty(remarks) ? tableName : remarks);

                ResultSet columnResultSet = databaseMetaData.getColumns(null, null, tableName, null);
                Map<String, String> columnCommentMap = new HashMap<>();
                while (columnResultSet.next()) {
                    columnCommentMap.put(columnResultSet.getString("COLUMN_NAME"), columnResultSet.getString("REMARKS"));
                }
                List<String> primaryColumnNames = new ArrayList<>();
                try (ResultSet rs = databaseMetaData.getPrimaryKeys(connection.getCatalog(), null, tableName)) {
                    while (rs.next()) {
                        String primaryKey = rs.getString("COLUMN_NAME");
                        primaryColumnNames.add(primaryKey);
                    }
                }

                List<ColumnMetadata> columnMetadataList = this.buildColumnMetadata(connection, tableName);
                columnMetadataList.forEach(c -> {
                    String colRemarks = columnCommentMap.get(c.getName());
                    c.setComment(StrUtils.isEmpty(colRemarks) ? c.getName() : colRemarks);
                    this.changeJavaType(globalConfig, c);
                    if (primaryColumnNames.contains(c.getName())) {
                        c.setPrimary(true);
                    }
                });

                tableMetadata.setColumns(columnMetadataList);
                tableMetadataList.add(tableMetadata);
            }

            return tableMetadataList;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private List<ColumnMetadata> buildColumnMetadata(Connection connection, String tableName) {
        try (Statement stm = connection.createStatement(); ResultSet rs = stm.executeQuery(String.format(SQL, tableName))) {

            ResultSetMetaData rsColumnMetaData = rs.getMetaData();
            int columnCount = rsColumnMetaData.getColumnCount();

            List<ColumnMetadata> columnMetadataList = new ArrayList<>();
            for (int i = 1; i <= columnCount; i++) {
                ColumnMetadata columnMetadata = new ColumnMetadata();
                columnMetadata.setName(rsColumnMetaData.getColumnName(i));
                columnMetadata.setColumnType(rsColumnMetaData.getColumnTypeName(i));
                columnMetadata.setLength(rsColumnMetaData.getColumnDisplaySize(i));
                String jdbcType = rsColumnMetaData.getColumnClassName(i);
                columnMetadata.setJavaType(ReflectUtils.forName(jdbcType));
                columnMetadata.setAutoIncrement(rsColumnMetaData.isAutoIncrement(i));
                columnMetadata.setNullable(rsColumnMetaData.isNullable(i) == ResultSetMetaData.columnNullable);

                columnMetadataList.add(columnMetadata);
            }
            return columnMetadataList;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

}
